#include "RBMainPCH.h"
#include "sh_math.h"

#include "DirectX/Include/D3DX10math.h"

namespace RB
{
namespace sh_math
{

float P(int l,int m,float x) 
{  
    // evaluate an Associated Legendre Polynomial P(l,m,x) at x 
    float pmm = 1.0f; 
    if(m>0) { 
        float somx2 = sqrt((1.0f-x)*(1.0f+x)); 

        float fact = 1.0f;
        for(int i=1; i<=m; i++) { 
            pmm *= (-fact) * somx2; 
            fact += 2.0f; 
        }
    }
    
    if(l==m) return pmm; 
    
    float pmmp1 = x * (2.0f*m+1.0f) * pmm; 
    if(l==m+1) return pmmp1; 

    float pll = 0.0f; 
    for(int ll=m+2; ll<=l; ++ll) 
    { 
        pll = ( (2.0f*ll-1.0f)*x*pmmp1-(ll+m-1.0f)*pmm ) / (ll-m); 
        pmm = pmmp1; 
        pmmp1 = pll; 
    }
    
    return pll; 
}


int factorial(int k)
{
    int res=1;
    while(k>0) 
    {
        res = res * k;
        --k;
    }
    return res;
}

float K(int l, int m) 
{
    // renormalisation constant for SH function 
    float temp = ((2.0f*l+1.0f)*factorial(l-m)) / (4.0f*RBMath::RB_PI*factorial(l+m)); 
    return sqrt(temp); 
}

float SH(int l, int m, float theta, float phi) 
{
     // return a point sample of a Spherical Harmonic basis function 
     // l is the band, range [0..N] 
     // m in the range [-l..l] 
     // theta in the range [0..Pi] 
     // phi in the range [0..2*Pi] 
     static const float sqrt2 = sqrt(2.0f); 
     if(m==0) return K(l,0)*P(l,0,cos(theta)); 
     else if(m>0) return sqrt2*K(l,m)*cos(m*phi)*P(l,m,cos(theta)); 
     else return sqrt2*K(l,-m)*sin(-m*phi)*P(l,-m,cos(theta)); 
}

void SH(float* sh, const float* direction, int order)
{
    float phi = atan(direction[1]/direction[0]);
    float theta = acos(direction[2]);
    float factor = 1.0f;
    for (int l=0; l<order; ++l)
    {
        for (int m=-l; m<=l; ++m)
        {
            int idx = l*l+m+l;
            // Check for sign change in m
            sh[idx] = SH(l, m, theta, phi);
        }
    }
}

float dot(const float* lsh, const float* rsh, int order)
{
    float res = 0;
    const int coeff_count = order*order;
    for (int i=0; i < coeff_count; ++i)
    {
        res += lsh[i] * rsh[i];
    }
    return res;
}


void eval_direction( const float* dir, float* sh, int order )
{
    float x = dir[0];
    float y = dir[1];
    float z = dir[2];

    switch (order)
    {
    default:
        D3DXSHEvalDirection(sh, order, (D3DXVECTOR3*)dir);
        break;
    case 3:
        sh[4] = 1.092548f * x*y;        // 2,-2
        sh[5] = -1.092548f * y*z;       // 2,-1
        sh[6] = 0.315392f * (3*z*z-1);  // 2, 0
        sh[7] = -1.092548f * x*z;       // 2, 1
        sh[8] = 0.546274f * (x*x-y*y);  // 2, 2
    case 2:
        sh[1] = -0.488603f * y;         // 1,-1
        sh[2] = 0.488603f * z;          // 1, 0
        sh[3] = -0.488603f * x;         // 1, 1
    case 1:
        sh[0] = 0.282095f;              // 0, 0
    }

}


float3 get_cube_map_direction(int i, int j, int face, int res, float* weight)
{
    float x = i - res * 0.5f + 0.5f;
    float y = j - res * 0.5f + 0.5f;
    float z = res *  0.5f;
    
    if (weight != NULL)
    {
        float x0 = (i - res * 0.5f) / z;
        float y0 = (j - res * 0.5f) / z;
        float x1 = (i - res * 0.5f + 1) / z;
        float y1 = (j - res * 0.5f + 1) / z;

        *weight = 
             atanf((x1 * y1) / sqrtf(1 + x1 * x1 + y1 * y1)) - 
             atanf((x0 * y1) / sqrtf(1 + x0 * x0 + y1 * y1)) - 
             atanf((x1 * y0) / sqrtf(1 + x1 * x1 + y0 * y0)) +
             atanf((x0 * y0) / sqrtf(1 + x0 * x0 + y0 * y0)); 

    }

    switch (face)
    {
    case 0 : return float3( z, -x, -y).Normalize();
    case 1 : return float3(-z, -x,  y).Normalize();
    case 2 : return float3( y,  z,  x).Normalize();
    case 3 : return float3( y, -z, -x).Normalize();
    case 4 : return float3( y, -x,  z).Normalize();
    case 5 : return float3(-y, -x, -z).Normalize();
    }

    return float3::Zero();
}

void project_cube_map(float* sh_r, float* sh_g, float* sh_b, const void* const* data, size_t pitch, size_t res, size_t pixel_size, int order, rgb_reader reader )
{
    float basis[max_order*max_order];
    float r,g,b;

    float factor=1.0f;// / (4*RBMath::RB_PI);

    const int coeff_count = order*order;

    for (size_t i = 0; i < res; i++)
    {
        for (size_t j = 0; j < res; j++)
        {
            float weight;
            for (size_t f = 0; f < 6; f++)
            {
                const char* pixel_data_face = (const char*)data[f];
                const char* pixel_value = (const char*)&pixel_data_face[i * pitch + j * pixel_size];

                reader((const void*)pixel_value, &r, &g, &b);

                float3 dir = get_cube_map_direction(i, j, f, res, (f==0) ? &weight : NULL);
                eval_direction(&dir.x, &basis[0], order);

                for (int32 coeff_idx=0; coeff_idx < coeff_count; ++coeff_idx)
                {
                    sh_r[coeff_idx] += basis[coeff_idx] * (factor * r * weight);
                    sh_g[coeff_idx] += basis[coeff_idx] * (factor * g * weight);
                    sh_b[coeff_idx] += basis[coeff_idx] * (factor * b * weight);
                }
            }
        }
    }
}

void unproject_cube_map_face( const float* sh_r, const float* sh_g, const float* sh_b, void* data, int face, size_t pitch, size_t res, size_t pixel_size, int order, rgb_writer writer )
{
    float basis[max_order*max_order];
    float r,g,b;

    const int coeff_count = order*order;

    for (size_t i = 0; i < res; i++)
    {
        for (size_t j = 0; j < res; j++)
        {
            char* pixel_data_face = (char*)data;
            char* pixel_value = (char*)&pixel_data_face[i * pitch + j * pixel_size];

            float3 dir = get_cube_map_direction(i, j, face, res, NULL);
            eval_direction(&dir.x, &basis[0], order);

            r = dot(sh_r, basis, order);
            g = dot(sh_g, basis, order);
            b = dot(sh_b, basis, order);

            writer((void*)pixel_value, r, g, b);
        }
    }
}

}
}